'''
Created on Sep 18, 2009

@author: mkiyer
'''

import veggie.coverage.covdb as covdb
from veggie.coverage.visualization.ucsc import get_bedgraph_track_header, gen_bedgraph_track_data
from veggie.io.bed import bed_reader, simple_bed_reader
from veggie.sample.sampledb2 import get_sampledb
from veggie.sample.samplegroup import parse_samplegroups_xml
from veggie.gene.gene import parse_refseq_genes
from optparse import OptionParser
import collections
import h5py
import logging
import matplotlib.pyplot as plt
import numpy as np
import veggie.coverage.visualization.plots as oncoplot
import sys


def plot_expr(chrom, start, end, sample_groups,
              **kwargs):
    '''
    kwargs:
    figsize - size of figure (tuple)
    title - title of plot (string)
    shade_start - instructs certain plots to shade a region
    shade_end - instructs certain plots to shade a region
    '''
    # get coverage data
    sample_coverages = covdb.calc_avg_cov_arrays(chrom, start, end, sample_groups)
    summed_coverages = {}
    for s, covarray in sample_coverages.iteritems():
        summed_coverages[s] = covdb.calc_rpkm_array(covarray)

    # setup figure
    default_suptitle = 'Expression for %s:%d-%d' % (chrom, start, end)
    default_figsize = (15,10)
    #default_window_title = 'Boxplot'
    fig = plt.figure(figsize=kwargs.get('figsize', default_figsize))
    fig.suptitle(kwargs.get('title', default_suptitle))
    #fig.canvas.set_window_title(kwargs.get('window_title', default_window_title))
    fig.subplots_adjust(left=0.05, right=0.975, top=0.95, bottom=0.125,
                        wspace=0.1, hspace=0.25)

    ax1 = fig.add_subplot(221)
    ax2 = fig.add_subplot(223)
    ax3 = fig.add_subplot(222)
    ax4 = fig.add_subplot(224)
    # Top plot shows individual samples
    oncoplot.plot_sample_coverage(chrom, start, end, 
                                  sample_groups, sample_coverages,
                                  axes=ax1,
                                  **kwargs)
    # Bottom plot shows average for groups
    oncoplot.plot_sample_group_average_coverage(chrom, start, end, 
                                                sample_groups, 
                                                sample_coverages,
                                                axes=ax2,
                                                **kwargs)
    # Barplot
    oncoplot.plot_barplot(chrom, start, end,
                          sample_groups, summed_coverages,
                          axes=ax3)
    # Boxplot
    oncoplot.plot_boxplot(chrom, start, end, 
                          sample_groups, summed_coverages,
                          axes=ax4)
    plt.savefig('%s_%d-%d_expr.png' % (chrom, start, end))
    #plt.show()
    plt.close()
    return

def write_sample_header(sample_groups, outfhd):
    # sample properties header
    outfhd.write('Sample Group')
    for sgroup in sample_groups:
        for i in xrange(len(sgroup)):
            outfhd.write('\t%s' % sgroup.name)
    outfhd.write('\n')
    outfhd.write('Sample Name')
    for sgroup in sample_groups:
        for sample_name in sgroup.iterkeys():
            outfhd.write('\t%s' % sample_name)
    outfhd.write('\n')
    # write all properties in param_schema
    for prop in sorted(get_sampledb().param_schema.keys()):
        outfhd.write('%s' % prop)
        for sgroup in sample_groups:
            for sample_name in sgroup.iterkeys():
                if prop in get_sampledb().params[sample_name]:
                    outfhd.write('\t%s' % get_sampledb().params[sample_name][prop])
                else:
                    outfhd.write('\tNone')
        outfhd.write('\n')
    outfhd.write('\n')

def write_data_header(sample_groups, outfhd):
    outfhd.write('interval')
    for sgroup in sample_groups:
        for sample_name in sgroup.iterkeys():
            outfhd.write('\t%s' % sample_name)
    outfhd.write('\n')

def write_expression_matrix(intervals, sample_groups, outfhd,
                            pool_lanes=True):
    # output the sample names, groups, and parameters
    write_sample_header(sample_groups, outfhd)
    write_data_header(sample_groups, outfhd)
    # get coverage at each interval
    for interval in intervals:
        chrom, start, end = interval
        outfhd.write('%s:%d-%d' % (chrom, start, end))
        
        for sgroup in sample_groups:
            if pool_lanes:
                for sample_dsets in sgroup.itervalues():
                    # here we have the interval, the sample name, and the
                    # datasets that comprise the sample.  we can now 
                    # compute coverage
                    avgcov = covdb.calc_avg_cov_rpkm(chrom, start, end, sample_dsets)
                    outfhd.write('\t%.2f' % avgcov)
            else:
                for sample_dsets in sgroup.itervalues():                    
                    rpkms = [('%.2f' % covdb.calc_rpkm(chrom, start, end, dset)) 
                             for dset in sample_dsets]
                    outfhd.write('\t')
                    outfhd.write(','.join(rpkms))
        outfhd.write('\n')

def write_bedgraph(intervals, sample_groups, outfhd):
    for sgroup in sample_groups:
        for sample_name, sample_dsets in sgroup.iteritems():
            outfhd.writelines(get_bedgraph_track_header(sample_name, sample_name))
            for interval in intervals:
                chrom, start, end = interval.chrom, interval.start, interval.end
                avgcovarr = 1000.0 * covdb.calc_avg_cov_array(chrom, start, end, sample_dsets)                
                lines = gen_bedgraph_track_data(chrom, start, end, avgcovarr)
                outfhd.writelines(lines)

def get_intervals_from_gene_list(fhd,
                                 exons=True,
                                 introns=False,
                                 isoforms=True):
    logging.debug("building gene lookup table")
    genes = {}
    for g in parse_refseq_genes():
        genes[g.symbol] = g
        genes[g.acc] = g
    intervals = []
    for thisline in open(args[0]):
        gene_name = thisline.strip()
        if gene_name is None:
            continue
        if gene_name.startswith('#'):
            continue
        g = genes.get(gene_name, None)
        if g is None:
            logging.debug('Could not find gene %s, skipping...\n' % gene_name)
            continue
        logging.debug('Gene %s: %s:%d-%d' % (gene_name, g.chrom, g.tx_start, g.tx_end))
        for e in g.exons:
            intervals.append((g.chrom, e[0], e[1]))
    return intervals
        
if __name__ == '__main__':
    # configure logging
    logging.basicConfig(level=logging.DEBUG)

    # parse command line
    optionparser = OptionParser("usage: %prog [options] <regions.bed>")
    optionparser.add_option("--covdb", dest="covdb",
                            help="path to coverage database")
    optionparser.add_option("-s", "--samples", dest="samples",
                            help="sample group XML file")
    optionparser.add_option("-o", "--output", dest="outfile",
                            default=None,
                            help="output file [default: %default]")
    optionparser.add_option("-m", "--matrix",
                            action="store_true", dest="matrix", default=False,
                            help="output expression matrix")
    optionparser.add_option("-p", "--plots",
                            action="store_true", dest="plots", default=False,
                            help="create plots for each interval")
    optionparser.add_option("--bedgraph",
                            action="store_true", dest="bedgraph", default=False,
                            help="generate bedgraph track for UCSC browser")
    optionparser.add_option("-g", action="store_true", 
                            dest="genes_input", default=False,
                            help="the region file contains a list of genes")
    optionparser.add_option("--lanes", action="store_false", dest="pool_lanes",
                            default=True,
                            help="show individual lane results rather "
                            "than pooling")

    (options, args) = optionparser.parse_args()
    if options.outfile == None:
        outfhd = sys.stdout
    else:
        outfhd = open(options.outfile, 'w')
    if len(args) == 0:
        optionparser.error("no BED file specified")

    # get list of intervals
    if options.genes_input:
        logging.error("gene list temporarily disabled")
        sys.exit(1)
        intervals = sorted(get_intervals_from_gene_list(open(args[0])))
    else:
        intervals = sorted(list(bed_reader(open(args[0]))))
    
    # open coverage db
    h5file = h5py.File(options.covdb, 'r')
    # generate SampleGroupCoverage objects that pair samples with
    # coverage data
    sample_groups = [covdb.SampleGroupCoverage(sgroup, h5file) for sgroup in 
                     list(parse_samplegroups_xml(options.samples))]

    if options.matrix:
        write_expression_matrix(intervals, sample_groups, outfhd,
                                pool_lanes=options.pool_lanes)

    if options.plots:
        # process intervals
        for interval in intervals:
            logging.debug('plotting %s:%d-%d' % (interval.chrom, interval.start, interval.end))
            title='%s:%d-%d name=%s strand=%s' % \
                (interval.chrom, interval.start, interval.end, interval.name, interval.strand)
            plot_expr(interval.chrom, interval.start, interval.end, sample_groups,
                      title=title,
                      shade_start=interval.thickStart)
    
    if options.bedgraph:
        write_bedgraph(intervals, sample_groups, outfhd)
    
    h5file.close()
